# import sqlite3
import os
import psycopg2
import exceptions

class DBConnection:
    # in order to trace the source of postgresql query in postgresql log, add environment variable 'PGAPPNAME' to specify the application name.   
    def __init__(self, _db_name, _user, _password, _app_name, _auto_commit=False):
        try:
            os.environ['PGAPPNAME'] = _app_name
            # Connect to db
            self._connection = psycopg2.connect(database=_db_name, user=_user, password=_password)
            if _auto_commit :
                self._connection.autocommit=True
            
            # Get cursor
            self._cursor = self._connection.cursor()
        except psycopg2.OperationalError:
            raise Exception()

    def __del__(self):
        # Disconnect from db.
        # We check if _connection attr exists because if db connect fail, the attribute is not set.
        if hasattr(self, "_connection"):
            self._connection.close()

    def execute(self, *args):
        return self._cursor.execute(*args)

    def commit(self):
        return self._connection.commit()

    def fetchone(self):
        return self._cursor.fetchone()

    def fetchall(self):
        return self._cursor.fetchall()

    def rowcount(self):
        return self._cursor.rowcount

    def fetch_column_value(self):
        columns = [ x[0] for x in self._cursor.description ]
        results = self.fetchall()
        retdata = [ dict(zip(columns, rst)) for rst in results ]
        return retdata

    def rollback(self):
        return self._connection.rollback()

